from __future__ import division
import os
import torch.nn.functional as F
import numpy as np
from utils.preprocessing import *
from utils.plotting import *
from utils.metrix import *


def predict(model, data_set, data_loader, do_blur=False, counting=False):

    model.eval()

    num_smps = len(data_set)
    num_kps = data_set.num_kps

    print num_smps, num_kps

    preds_xy = np.empty((num_smps, num_kps, 2))  # (N,24,2) x y
    true_xy = np.empty((num_smps, num_kps, 3))  # (N,24,3) x y vis
    all_ori_sizes = np.empty((num_smps, 2))   # (N, 2)
    all_true_flms = np.empty((num_smps, num_kps, 3))
    all_cates = np.empty(num_smps)

    idx = 0

    for bc_cnt, bc_data in enumerate(data_loader):
        if counting:
            print('%d/%d' % (bc_cnt, len(data_set)//data_loader.batch_size))

        imgs, lm_masks, vis_masks, ori_sizes, flms, cate_idxs = bc_data

        imgs = Variable(imgs.cuda())

        outputs = F.sigmoid(model(imgs))

        # post process and find peak points   (bs, C, 2)  x,y
        batch_peaks = batch_postprocess(outputs.data.cpu().numpy(), ori_sizes, do_blur=do_blur)

        # seat bc result to placeholder of the entire dataset
        preds_xy[idx: idx + imgs.size(0)] = batch_peaks
        all_ori_sizes[idx: idx + imgs.size(0)] = ori_sizes
        all_true_flms[idx: idx + imgs.size(0)] = flms
        all_cates[idx: idx + imgs.size(0)] = cate_idxs

        # never forget to update idx
        idx += imgs.size(0)

    # get true xy
    # (there might be tiny difference between ori true xy and the one computed by flm)
    true_xy[:, :, 0] = all_true_flms[:, :, 0] * all_ori_sizes[:, 1][:, np.newaxis]  # ori x coord
    true_xy[:, :, 1] = all_true_flms[:, :, 1] * all_ori_sizes[:, 0][:, np.newaxis]  # ori y coord
    true_xy[:, :, 2] = all_true_flms[:, :, 2]

    return preds_xy.round().astype(int),  true_xy.round().astype(int),  all_cates
